#include "stream_pusher.h"

namespace stream {

HBStreamPusher::HBStreamPusher(StreamPusherInfo* pusher_ifo, LogInfo* log_info): 
  BasePusher(pusher_ifo, log_info) {

}

HBStreamPusher::~HBStreamPusher() {

}
int HBStreamPusher::global_init() {
  int ret = 0;
  ret = HB_VENC_Module_Init();
  if (ret != 0) { spdlog::get("logger")->info("HB_VENC_Module_Init: {}\n", ret); }
  return ret;
}
int HBStreamPusher::global_release(){
  int ret = HB_VENC_Module_Uninit();
  if (ret != 0) { printf("HB_VENC_Module_Init: %d\n", ret); }
}

void HBStreamPusher::release_ffmpeg() {
  // TODO:详细研究释放规则
  // avformat_close_input(&pfmt_ctx);
  // avcodec_close(out_codec_ctx);
    // avcodec_free_context(out_codec_ctx);
  avformat_free_context(pfmt_ctx);
  pfmt_ctx = nullptr;
}

void HBStreamPusher::release_encoder() {
  int ret = 0;
  ret = HB_VENC_StopRecvFrame(pusher_ifo->vEncChn);
  if (ret != 0) { printf("HB_VENC_StopRecvFrame failed\n"); }

  ret = HB_VENC_DestroyChn(pusher_ifo->vEncChn);
  if (ret != 0) { printf("HB_VENC_DestroyChn failed\n"); }

}

int HBStreamPusher::init_pusher_params() {
  int ret = 0;
  char errbuf[64];

  av_register_all();       // Register all codecs and formats so that they can be used.
  ret = avformat_network_init();
  if ((ret = avformat_alloc_output_context2(&pfmt_ctx, nullptr, pusher_ifo->formatName.c_str(), nullptr)) < 0) {
    std::cout << "Could not allocate output format context!" << std::endl;
    return -1;
  }

  if (!(pfmt_ctx->oformat->flags & AVFMT_NOFILE)) {
    if ((ret = avio_open2(&pfmt_ctx->pb, pusher_ifo->streamAddress.c_str(), AVIO_FLAG_WRITE, nullptr, nullptr)) < 0) {
    std::cout << "Could not open output IO context!" << std::endl;
    return -1;
    }
  }

  AVCodecID encoder_type = AV_CODEC_ID_NONE;
  if (pusher_ifo->enType == AV_CODEC_ID_H264) { 
    encode_type = PT_H264; 
    encoder_type = AV_CODEC_ID_H264; 
  }
  else if (pusher_ifo->enType == AV_CODEC_ID_HEVC) { 
    encode_type = PT_H265; 
    encoder_type = AV_CODEC_ID_HEVC; 
  }
  out_codec = avcodec_find_encoder(encoder_type);
  out_stream = avformat_new_stream(pfmt_ctx, out_codec);
  out_codec_ctx = avcodec_alloc_context3(out_codec);

  out_codec_ctx->codec_tag = 0;
  out_codec_ctx->codec_id = out_codec->id;
  out_codec_ctx->codec_type = AVMEDIA_TYPE_VIDEO;
  out_codec_ctx->width = pusher_ifo->dstWidth;
  out_codec_ctx->height = pusher_ifo->dstHeight;
  out_codec_ctx->gop_size = pusher_ifo->fps;
  out_codec_ctx->pix_fmt = AV_PIX_FMT_NV12;
  out_codec_ctx->framerate = {pusher_ifo->fps, 1};
  out_codec_ctx->time_base = {1, pusher_ifo->fps};
  out_codec_ctx->bit_rate = pusher_ifo->bitrate * 1000;
  if (pfmt_ctx->oformat->flags & AVFMT_GLOBALHEADER) {
    out_codec_ctx->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;
  }

  out_codec_ctx->thread_count = 3;
  out_codec_ctx->max_b_frames = 0;  //编码器不使用Ｂ帧，提高直播解码效率


  if ((ret = avcodec_parameters_from_context(out_stream->codecpar, out_codec_ctx)) < 0) {
    std::cout << "Could not initialize stream codec parameters!" << std::endl;
    return -1;
  }

  
  AVDictionary *codec_options = nullptr;
  av_dict_set(&codec_options, "profile", pusher_ifo->codecProfile.c_str(), 0);
  av_dict_set(&codec_options, "preset", "superfast", 0);
  av_dict_set(&codec_options, "tune", "zerolatency", 0);
  // //无需打开编码器
  // open video encoder
  if ((ret = avcodec_open2(out_codec_ctx, out_codec, &codec_options)) < 0) {
    std::cout << "Could not open video encoder!" << std::endl;
    return -1;
  }
  av_dict_free(&codec_options);


  out_stream->codecpar->extradata = out_codec_ctx->extradata;
  out_stream->codecpar->extradata_size = out_codec_ctx->extradata_size;

  out_stream->codecpar->codec_tag = 0;
  av_dump_format(pfmt_ctx, 0, pusher_ifo->streamAddress.c_str(), 1);

  if ((ret = avformat_write_header(pfmt_ctx, nullptr)) < 0) {
    std::cout << "Could not write header!" << std::endl;
    return -1;
  }
  
  return 0;
}

int HBStreamPusher::add_frame(stream::ImageBlob& buffer) {
  int ret = 0;
  queue_lock.lock();
  if (queue_data.size() > 5) { 
    queue_data.pop(); 
    printf("HBStreamPusher %d pop.\n", pusher_ifo->vEncChn);
  }
  queue_data.push(buffer);
  // std::cout << "HBStreamPusher::add_frame queue_data: " << queue_data.size() << std::endl;
  queue_lock.unlock();
  return 0;
}

int HBStreamPusher::init_venc_params() {
  int ret = 0;
  VENC_CHN_ATTR_S vencChnAttr;
  // 初始化channel属性
  ret = venc_ChnAttr_init(&vencChnAttr);
  if (ret != 0) { printf("sample_venc_ChnAttr_init failded: %d\n", ret); }
  // 创建channel
  ret = HB_VENC_CreateChn(pusher_ifo->vEncChn, &vencChnAttr);
  if (ret != 0) {
    printf("HB_VENC_CreateChn %d failed, %d.\n", pusher_ifo->vEncChn, ret);
    return -1;
  }
  // 配置Rc参数
  ret = venc_setRcParam(pusher_ifo->vEncChn, pusher_ifo->bitrate, &vencChnAttr);
  if (ret != 0) { printf("sample_venc_setRcParam failded: %d\n", ret); }
  // 设置channel属性
  ret = HB_VENC_SetChnAttr(pusher_ifo->vEncChn, &vencChnAttr);  // config
  if (ret != 0) {
    printf("HB_VENC_SetChnAttr failed\n");
    return -1;
  }
  VENC_RECV_PIC_PARAM_S pstRecvParam;
  pstRecvParam.s32RecvPicNum = 0;  // unchangable
  ret = HB_VENC_StartRecvFrame(pusher_ifo->vEncChn, &pstRecvParam);
  if (ret != 0) {
    printf("HB_VENC_StartRecvFrame failed\n");
    return -1;
  }
  return 0;
}
int HBStreamPusher::start_push_listener() {
  // stop_flag2 = common::CameraGrabing;
  int ret = start_push_imp();
  printf("HBStreamPusher::start_push_listener start_push_imp:%d.\n", ret);
  // stop_flag2 = common::CameraOpened;
  return  ret;
}
int HBStreamPusher::start_push_imp() {
  int ret = 0;

  int pts = 0;
  VIDEO_FRAME_S pstFrame;
  hb_vio_buffer_t vio_buf;
  memset(&vio_buf, 0, sizeof(hb_vio_buffer_t));
  int size_y = pusher_ifo->dstWidth * pusher_ifo->dstHeight;
  int size_uv = size_y / 2;
  ret = prepare_user_buf(&vio_buf, size_y, size_uv);
  if (ret != 0) {printf("error. prepare_user_buf ret:%d.\n", ret); return ret;}
  
  vio_buf.img_info.planeCount = 2;
  vio_buf.img_info.img_format = 8;
  vio_buf.img_addr.width = pusher_ifo->dstWidth;
  vio_buf.img_addr.height = pusher_ifo->dstHeight;
  vio_buf.img_addr.stride_size = pusher_ifo->dstWidth;
  printf("HBStreamPusher::start_push_imp. init_params ok vEncChn:%d.\n", pusher_ifo->vEncChn);

  float venc_fps;
  int pts_count = 0, fps_cnt = 0;
  int64_t start_time = 0, end_time = 0;
  stream::ImageBlob inBuffer(stream::ImageBlobMode_YUV);
  int for_sleep = 5;
  stop_flag2 = common::CameraGrabing;
  while (stop_flag2 == common::CameraGrabing) {
    std::this_thread::sleep_for(std::chrono::milliseconds(for_sleep));

    if (fps_cnt == 0) { start_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();}
        
    queue_lock.lock();
    if (queue_data.empty()) { queue_lock.unlock(); continue; }
    inBuffer = queue_data.front();
    queue_data.pop();
    queue_lock.unlock();

    fps_cnt++;
    std::copy(inBuffer.yuv.begin(), inBuffer.yuv.begin()+size_y, vio_buf.img_addr.addr[0]);
    std::copy(inBuffer.yuv.begin()+size_y, inBuffer.yuv.end(), vio_buf.img_addr.addr[1]);

    pstFrame.stVFrame.width = pusher_ifo->dstWidth;
    pstFrame.stVFrame.height = pusher_ifo->dstHeight;
    pstFrame.stVFrame.size = size_y + size_uv;
    pstFrame.stVFrame.pix_format = HB_PIXEL_FORMAT_NV12;
    pstFrame.stVFrame.phy_ptr[0] = vio_buf.img_addr.paddr[0];
    pstFrame.stVFrame.phy_ptr[1] = vio_buf.img_addr.paddr[1];
    pstFrame.stVFrame.vir_ptr[0] = vio_buf.img_addr.addr[0];
    pstFrame.stVFrame.vir_ptr[1] = vio_buf.img_addr.addr[1];
    pstFrame.stVFrame.pts = pts;
    
    ret = HB_VENC_SendFrame(pusher_ifo->vEncChn, &pstFrame, 3000);
    end_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    if (pts % 300 == 0) { printf("venc SendFrame %d, ret=%d, %lld\n", pts, ret, end_time); }
    pts++;

    if (ret != 0) { continue; }

    // 输出fps
    if (end_time - start_time > 5 * 1e3) {
      venc_fps = 1.0 / ((end_time - start_time) / (fps_cnt * 1e3));
      if (log_ifo->log_level_2) spdlog::get("logger")->info("chn:{}, venc_fps:{}", pusher_ifo->vEncChn, venc_fps);
      printf("venc_fps c:%d : %f\n", pusher_ifo->vEncChn, venc_fps);
      fps_cnt = 0;
    }
  }

  stop_flag2 = common::CameraClosed;
  return 0;
}

int HBStreamPusher::init_params() {
  int ret = 0;
  if ((ret = init_pusher_params()) != 0) {
    printf("[ERROR] init_pusher_params %d error\n", pusher_ifo->vEncChn);
    return ret;
  }

  if ((ret = init_venc_params()) != 0) {
    printf("[ERROR] init_venc_params %d error\n", pusher_ifo->vEncChn);
    return ret;
  }
  return 0;
}
int HBStreamPusher::start_push() {
  int ret = 0;

  // 初始化完成
  std::thread(&HBStreamPusher::start_push_listener, this).detach();

  int frame_cnt = 0;
  int64_t pts = 0;
  float push_fps;
  int pts_count = 0, fps_cnt = 0;
  int64_t start_time = 0, end_time = 0;
  int64_t prev_time, cur_time;
  prev_time = cur_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();;
  VIDEO_STREAM_S pstStream;
  int step_time = 1000 / pusher_ifo->fps;
  int send_ok = 0;
  // AVBitStreamFilterContext h264bsfc = av_bsf_init();
  stop_flag = common::CameraGrabing;
  while (stop_flag == common::CameraGrabing) {
    cur_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    if (fps_cnt == 0) { start_time = cur_time; }

    // 控制发送速度
    if (send_ok >= 0) {
      int leave_time = step_time - (cur_time - prev_time);
      if (leave_time > 0) std::this_thread::sleep_for(std::chrono::milliseconds(leave_time));
      prev_time = cur_time + leave_time;
    } else {
      std::this_thread::sleep_for(std::chrono::milliseconds(2));
    }

    ret = HB_VENC_GetStream(pusher_ifo->vEncChn, &pstStream, 2000);
    // printf("venc GetStream %d ret=%d\n", pts, ret);
    if (ret != 0) { send_ok = -1; continue; }
    // fwrite(pstStream.pstPack.vir_ptr,
    //                 pstStream.pstPack.size, 1, outFile);


    fps_cnt++;
    AVPacket *encode_pkt = av_packet_alloc();
    ret = av_new_packet(encode_pkt, pstStream.pstPack.size);
    memcpy(encode_pkt->data, pstStream.pstPack.vir_ptr, pstStream.pstPack.size);
    

    pts += av_rescale_q(1, out_codec_ctx->time_base, out_stream->time_base);
    encode_pkt->pts = pts;
    av_interleaved_write_frame(pfmt_ctx, encode_pkt);
    av_packet_unref(encode_pkt);
    av_packet_free(&encode_pkt);
    send_ok = 0;

    // int64_t cur_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();;
    // std::cout << frame_cnt++ << ".  encoder time: " << time2 - time1 << ".  total time: " << cur_time - timepoint << std::endl;
    // timepoint = cur_time;
        
    ret = HB_VENC_ReleaseStream(pusher_ifo->vEncChn, &pstStream);
    if (ret != 0) { printf("HB_VENC_ReleaseStream %d failed\n", pusher_ifo->vEncChn); }
    // std::cout << "stream_frame_size: " << stream_frame_size << std::endl;
    // int64_t time2 = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();;

    // 输出fps
    end_time = duration_cast<milliseconds>(system_clock::now().time_since_epoch()).count();
    if (end_time - start_time > 5 * 1e3) {
      push_fps = 1.0 / ((end_time - start_time) / (fps_cnt * 1e3));
      // spdlog::get("logger")->info("pull_fps: {}", pull_fps);
      printf("push_fps %d: %f\n", pusher_ifo->vEncChn, push_fps);
      fps_cnt = 0;
    }
  }

  release_ffmpeg();
  release_encoder();
  stop_flag = common::CameraClosed;
  return 0;

}


int HBStreamPusher::venc_setRcParam(int chn, int bitRate, VENC_CHN_ATTR_S *pVencChnAttr) {
  VENC_RC_ATTR_S *pstRcParam;
  int s32Ret;

  if (pVencChnAttr->stVencAttr.enType == PT_H264) {
    pstRcParam = &(pVencChnAttr->stRcAttr);
    // 为什么之前是VBR 这里改为CBR 之前的VBR是必须的吗？
    pVencChnAttr->stRcAttr.enRcMode = VENC_RC_MODE_H264CBR;
    s32Ret = HB_VENC_GetRcParam(chn, pstRcParam);
    if (s32Ret != 0) {
      printf("HB_VENC_GetRcParam failed.\n");
      return -1;
    }

    printf(" vencChnAttr.stRcAttr.enRcMode = %d mmmmmmmmmmmmmmmmmm   \n",
            pVencChnAttr->stRcAttr.enRcMode);
    printf(" u32VbvBufferSize = %d mmmmmmmmmmmmmmmmmm   \n",
            pVencChnAttr->stRcAttr.stH264Cbr.u32VbvBufferSize);
    // 设置码率
    pstRcParam->stH264Cbr.u32BitRate = bitRate;
    // 设置帧率
    pstRcParam->stH264Cbr.u32FrameRate = pusher_ifo->fps;
    // 设置I帧间隔
    pstRcParam->stH264Cbr.u32IntraPeriod = pusher_ifo->fps;
    // 设置VbvBufferSize，与码率、解码器速率有关，一般在1000～5000
    pstRcParam->stH264Cbr.u32VbvBufferSize = 3000;
  } else if (pVencChnAttr->stVencAttr.enType == PT_H265) {
    pstRcParam = &(pVencChnAttr->stRcAttr);
    pVencChnAttr->stRcAttr.enRcMode = VENC_RC_MODE_H265CBR;
    s32Ret = HB_VENC_GetRcParam(chn, pstRcParam);
    if (s32Ret != 0) {
      printf("HB_VENC_GetRcParam failed.\n");
      return -1;
    }
    printf(" m_VencChnAttr.stRcAttr.enRcMode = %d mmmmmmmmmmmmmmmmmm   \n",
            pVencChnAttr->stRcAttr.enRcMode);
    printf(" u32VbvBufferSize = %d mmmmmmmmmmmmmmmmmm   \n",
            pVencChnAttr->stRcAttr.stH265Cbr.u32VbvBufferSize);
    // 设置码率
    pstRcParam->stH265Cbr.u32BitRate = bitRate;
    // 设置帧率
    pstRcParam->stH265Cbr.u32FrameRate = pusher_ifo->fps;
    // 设置I帧间隔
    pstRcParam->stH265Cbr.u32IntraPeriod = pusher_ifo->fps;
    // 设置VbvBufferSize，与码率、解码器速率有关，一般在1000～5000
    pstRcParam->stH265Cbr.u32VbvBufferSize = 3000;
  }
  return 0;
}

int HBStreamPusher::venc_ChnAttr_init(VENC_CHN_ATTR_S *pVencChnAttr) {
  int streambufSize = 0;
  if (pVencChnAttr == NULL) {
    printf("pVencChnAttr is NULL!\n");
    return -1;
  }
  // 该步骤必不可少
  memset(pVencChnAttr, 0, sizeof(VENC_CHN_ATTR_S));
  // 设置编码模型分别为 PT_H264 PT_H265 PT_MJPEG 
  pVencChnAttr->stVencAttr.enType = (PAYLOAD_TYPE_E)encode_type;
  // 设置编码分辨率
  pVencChnAttr->stVencAttr.u32PicWidth = pusher_ifo->dstWidth;
  pVencChnAttr->stVencAttr.u32PicHeight = pusher_ifo->dstHeight;
  // 设置像素格式 NV12格式
  pVencChnAttr->stVencAttr.enPixelFormat = HB_PIXEL_FORMAT_NV12;
  // 配置图像镜像属性 无镜像
  pVencChnAttr->stVencAttr.enMirrorFlip = DIRECTION_NONE;
  // 配置图像旋转属性 不旋转
  pVencChnAttr->stVencAttr.enRotation = CODEC_ROTATION_0;
  // 配置图像剪裁属性 不剪裁
  pVencChnAttr->stVencAttr.stCropCfg.bEnable = HB_FALSE;
  // 输入图像大小 1024对齐
  int vlc_buf_size = pusher_ifo->dstWidth * pusher_ifo->dstHeight;
  streambufSize = (vlc_buf_size * 3 / 2 + 1024) & ~0x3ff;
  // vlc_buf_size 为经验值，可以减少RAM使用，如果想使用默认值则保持为0
  if (vlc_buf_size > 2688 * 1522) {
    pVencChnAttr->stVencAttr.vlc_buf_size = 7900 * 1024;
  } else if (vlc_buf_size > 1920 * 1080) {
    pVencChnAttr->stVencAttr.vlc_buf_size = 4 * 1024 * 1024;
  } else if (vlc_buf_size > 1280 * 720) {
    pVencChnAttr->stVencAttr.vlc_buf_size = 2100 * 1024;
  } else if (vlc_buf_size > 704 * 576) {
    pVencChnAttr->stVencAttr.vlc_buf_size = 2100 * 1024;
  } else {
    pVencChnAttr->stVencAttr.vlc_buf_size = 2048 * 1024;
  }
  if (pusher_ifo->enType == AV_CODEC_ID_JPEG2000 || pusher_ifo->enType == AV_CODEC_ID_MJPEG) {
    // 输出码流个数
    pVencChnAttr->stVencAttr.u32BitStreamBufferCount = 1;
    // 输入图像个数
    pVencChnAttr->stVencAttr.u32FrameBufferCount = 2;
    pVencChnAttr->stVencAttr.bExternalFreamBuffer = HB_TRUE;
    // 关闭DCF
    pVencChnAttr->stVencAttr.stAttrJpeg.dcf_enable = HB_FALSE;
    // 质量系数，越小质量越好
    pVencChnAttr->stVencAttr.stAttrJpeg.quality_factor = 0;
    // 配置为0 不建议更改
    pVencChnAttr->stVencAttr.stAttrJpeg.restart_interval = 0;
    // 4096对齐
    pVencChnAttr->stVencAttr.u32BitStreamBufSize = (streambufSize + 4096) & ~4095;
  } else {
    // 输出码流个数
    pVencChnAttr->stVencAttr.u32BitStreamBufferCount = 3;
    // 输入图像个数
    pVencChnAttr->stVencAttr.u32FrameBufferCount = 3;
    pVencChnAttr->stVencAttr.bExternalFreamBuffer = HB_TRUE;
    pVencChnAttr->stVencAttr.u32BitStreamBufSize = streambufSize;
  }

  if (pusher_ifo->enType == AV_CODEC_ID_H265) {
    // 配置编码模式为H265VBR
    pVencChnAttr->stRcAttr.enRcMode = VENC_RC_MODE_H265VBR;
    // 不使能QpMap
    pVencChnAttr->stRcAttr.stH265Vbr.bQpMapEnable = HB_FALSE;
    // 设置I帧Qp值
    pVencChnAttr->stRcAttr.stH265Vbr.u32IntraQp = 20;
    // 设置I帧间隔
    pVencChnAttr->stRcAttr.stH265Vbr.u32IntraPeriod = 60;
    // 设置帧率
    pVencChnAttr->stRcAttr.stH265Vbr.u32FrameRate = 30;
  }
  if (pusher_ifo->enType == AV_CODEC_ID_H264) {
    // 配置编码模式为H264VBR
    pVencChnAttr->stRcAttr.enRcMode = VENC_RC_MODE_H264VBR;
    // 不使能QpMap
    pVencChnAttr->stRcAttr.stH264Vbr.bQpMapEnable = HB_FALSE;
    // 设置I帧Qp值
    pVencChnAttr->stRcAttr.stH264Vbr.u32IntraQp = 20;
    // 设置I帧间隔
    pVencChnAttr->stRcAttr.stH264Vbr.u32IntraPeriod = pusher_ifo->fps;
    // 设置帧率
    pVencChnAttr->stRcAttr.stH264Vbr.u32FrameRate = pusher_ifo->fps;
    // h264_profile为0，系统自动配置
    pVencChnAttr->stVencAttr.stAttrH264.h264_profile = HB_H264_PROFILE_UNSPECIFIED;
    // h264_level为0，系统自动配置
    pVencChnAttr->stVencAttr.stAttrH264.h264_level = HB_H264_LEVEL_UNSPECIFIED;
  }
  // 设置GOP结构
  pVencChnAttr->stGopAttr.u32GopPresetIdx = 2;
  pVencChnAttr->stGopAttr.stCustomGopParam.u32CustomGopSize = pusher_ifo->fps;
  // 设置IDR帧类型
  pVencChnAttr->stGopAttr.s32DecodingRefreshType = 2;
  return 0;
}

int HBStreamPusher::stop_push() {
  int try_cnt = 200;
  if (stop_flag2 == common::CameraGrabing) {
    stop_flag2 = common::CameraOpened;
    while (stop_flag2 != common::CameraClosed && try_cnt-- > 0) {
      spdlog::get("logger")->info("HBStreamPusher::stop_push stop_flag2: {} ", stop_flag2);
      std::this_thread::sleep_for(milliseconds(50));
    }
    if (try_cnt < 0) {exit(0);}
  }
  try_cnt = 200;
  if (stop_flag == common::CameraGrabing) {
    stop_flag = common::CameraOpened;
    while (stop_flag != common::CameraClosed && try_cnt-- > 0) {
      spdlog::get("logger")->info("HBStreamPusher::stop_push stop_flag: {} ", stop_flag);
      std::this_thread::sleep_for(milliseconds(50));
    }
    if (try_cnt < 0) {exit(0);}
  }

  return 0;
}


}